A cornerstone of this method will be the powering of the diffusion matrix, which must intersperse jumping between points in the dataset and in the surrounding grid. Traditional matrix powering will do, provided we add the grid points to the dataset before powering, and then take them out afterwards.

A = torch.tensor([[1, 2], [3, 4], [5, 6]])
B = torch.rand(5, 2)
K = GaussianVectorField(2, 25, device=torch.device("cpu"))
out = diffusion_matrix_with_grid_points(A, B, K, 1, 0.5, 4)
out

tensor([[9.9995e-01, 4.7606e-05, 1.6631e-07],
        [2.2208e-03, 9.9237e-01, 5.4120e-03],
        [1.2162e-05, 3.4813e-03, 9.9651e-01]], grad_fn=<DivBackward0>)

General embedding matrix

The architecture should support modular swapping of point embedding methods. A few of these include:

Feedforward embedder (paired with a decoder)
Learnable initialized points (random, or to a prior embedding initialization, like DMaps, UMAP, PHATE...)
Fixed initializations, if we want the flow to contort itself to the initialization This will take the shape of swappable "embedder" classes, which can be supplied as input to the network.

For example, here is the simple FF net we use in the autoencoder:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

fa = FeedForwardReLU(shape=[2, 3, 4, 5, 6, 5, 4, 3, 2])
fa

FeedForwardReLU(
  (FA): Sequential(
    (0): Linear(in_features=2, out_features=3, bias=True)
    (1): LeakyReLU(negative_slope=0.01)
    (2): Linear(in_features=3, out_features=4, bias=True)
    (3): LeakyReLU(negative_slope=0.01)
    (4): Linear(in_features=4, out_features=5, bias=True)
    (5): LeakyReLU(negative_slope=0.01)
    (6): Linear(in_features=5, out_features=6, bias=True)
    (7): LeakyReLU(negative_slope=0.01)
    (8): Linear(in_features=6, out_features=5, bias=True)
    (9): LeakyReLU(negative_slope=0.01)
    (10): Linear(in_features=5, out_features=4, bias=True)
    (11): LeakyReLU(negative_slope=0.01)
    (12): Linear(in_features=4, out_features=3, bias=True)
    (13): LeakyReLU(negative_slope=0.01)
    (14): Linear(in_features=3, out_features=2, bias=True)
  )
)

%%javascript

IPython.tab_as_tab_everywhere = function(use_tabs) {
    if (use_tabs === undefined) {
        use_tabs = true; 
    }

    // apply setting to all current CodeMirror instances
    IPython.notebook.get_cells().map(
        function(c) {  return c.code_mirror.options.indentWithTabs=use_tabs;  }
    );
    // make sure new CodeMirror instances created in the future also use this setting
    CodeMirror.defaults.indentWithTabs=use_tabs;

    };

IPython.tab_as_tab_everywhere()

The Flow Embedder

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Basic Tests/Sanity Checks

More comprehensive explorations will be in the notebooks under 05c. For now, we ensure the network works.

from directed_graphs.datasets import (
    directed_circle,
    directed_cylinder,
    directed_spiral,
    directed_swiss_roll,
    directed_spiral_uniform,
    directed_swiss_roll_uniform,
)
from directed_graphs.datasets import plot_directed_2d, plot_directed_3d
from directed_graphs.diffusion_flow_embedding import DiffusionFlowEmbedder
import torch
import numpy as np

X, flow, labels = directed_circle(num_nodes=300, radius=1)
plot_directed_2d(X, flow, labels)
X = torch.tensor(X).float().to(device)
flow = torch.tensor(flow).float().to(device)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Input In [2], in <cell line: 16>()
     14 X, flow, labels = directed_circle(num_nodes=300, radius=1)
     15 plot_directed_2d(X, flow, labels)
---> 16 X = torch.tensor(X).float().to(device)
     17 flow = torch.tensor(flow).float().to(device)

NameError: name 'device' is not defined

MFE = MultiscaleDiffusionFlowEmbedder(X, flow, device=device).to(device)

MFE.fit(n_steps=1000)

  0%|          | 0/1000 [00:00<?, ?it/s]

Diffusion loss 0 is 2.03975510597229
Diffusion loss 1 is 2.73725962638855
Diffusion loss 2 is 3.0151374340057373
Diffusion loss 3 is 3.1220245361328125
EPOCH 0. Loss 3.1220245361328125. Flow strength 5.0. Heatmap of P embedding is

  0%|          | 2/1000 [00:01<11:27,  1.45it/s]

Diffusion loss 0 is 2.0292611122131348
Diffusion loss 1 is 2.730095863342285
Diffusion loss 2 is 3.0079665184020996
Diffusion loss 3 is 3.1147098541259766

  0%|          | 3/1000 [00:01<09:50,  1.69it/s]

Diffusion loss 0 is 2.0189754962921143
Diffusion loss 1 is 2.7233328819274902
Diffusion loss 2 is 3.0012006759643555
Diffusion loss 3 is 3.107822895050049

  0%|          | 4/1000 [00:02<09:01,  1.84it/s]

Diffusion loss 0 is 2.012047052383423
Diffusion loss 1 is 2.718916893005371
Diffusion loss 2 is 2.996764898300171
Diffusion loss 3 is 3.1032729148864746

  0%|          | 5/1000 [00:02<08:37,  1.92it/s]

Diffusion loss 0 is 2.0070102214813232
Diffusion loss 1 is 2.7159461975097656
Diffusion loss 2 is 2.993769884109497
Diffusion loss 3 is 3.1001739501953125

MFE.visualize_points(labels=labels)

No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.

plt.plot(MFE.losses["diffusion"], label="diffusion")
plt.plot(MFE.losses["smoothness"], label="smoothness")
plt.plot(MFE.losses["reconstruction"])
plt.legend()

<matplotlib.legend.Legend at 0x2b1bf09dfee0>

X, flow, labels = directed_swiss_roll_uniform(
    num_nodes=1000, num_spirals=2.5, radius=10, height=3, xtilt=0, ytilt=0
)
plot_directed_3d(X, flow, labels, mask_prob=0.5)

X = torch.tensor(X)
flow = torch.tensor(flow)
X = X.float().to(device)
flow = flow.float().to(device)

MFE = MultiscaleDiffusionFlowEmbedder(X, flow, device=device).to(device)

MFE.fit(n_steps=10000)

  0%|          | 0/10000 [00:00<?, ?it/s]

EPOCH 0. Loss 11.891934394836426. Flow strength 4.999000072479248. Heatmap of P embedding is

  1%|          | 100/10000 [00:15<19:12,  8.59it/s]

EPOCH 100. Loss 8.272184371948242. Flow strength 4.952264785766602. Heatmap of P embedding is

  2%|▏         | 200/10000 [00:30<26:09,  6.24it/s]

EPOCH 200. Loss 8.326654434204102. Flow strength 4.955779552459717. Heatmap of P embedding is

  3%|▎         | 300/10000 [00:44<22:10,  7.29it/s]

EPOCH 300. Loss 8.010849952697754. Flow strength 4.956110954284668. Heatmap of P embedding is

  4%|▍         | 400/10000 [00:59<22:36,  7.08it/s]

EPOCH 400. Loss 7.801275730133057. Flow strength 4.9596638679504395. Heatmap of P embedding is

!nbdev_build_lib

Converted 00_core.ipynb.
Converted 01a Graph Utils.ipynb.
Converted 01d Batch Testing Scripts.ipynb.
Converted 01e01 JSON Testing Script.ipynb.
Converted 02d01_Snake_TestCase.ipynb.
Converted 02d01a_Snake_With5b.ipynb.
Converted 02d02 Circles and Swiss roll Datasets.ipynb.
Converted 02d02a Testing Diffusion Flow Embedder on Circle and Swiss Roll.ipynb.
Converted 02d03 Spheres and Donuts with Flow (datasets).ipynb.
Converted 02d04 Tree and Flow Artist Test.ipynb.
Converted 02d05 Branch and Clusters.ipynb.
Converted 02d05a Testing Branch Clusters Circle Swiss roll.ipynb.
Converted 04a UW on RNA Data.ipynb.
Converted 05b Flow Embedding with Diffusion.ipynb.
Converted 05b01 Testing diffusion flow embedder.ipynb.
Converted 05b01a Diffusion Flow Embedding on Pancreas_RNA_Velocity.ipynb.
Converted 05b01a01 RNA Velocity Analysis.ipynb.
Converted 05b01a02 rna-velo_viz_phate.ipynb.
Converted 05b01a03 More_RNA_Velocity_Datasets.ipynb.
Converted 05b02 Flow-Artist_Upgrade.ipynb.
Converted 05b03 Flow-Artist-Scratch-Upgrade.ipynb.
Converted 05b03 Testing the Flow-Affinity Matrix.ipynb.
Converted 05b04 Diffusion Flow Embedder on Toy Data.ipynb.
Converted 05b05 Test Swiss Roll with New Kernel (7-11).ipynb.
Converted 05b05a DFE on the Swiss Roll with Smoothness.ipynb.
Converted 05c Multiscale Diffusion Flow Embedding.ipynb.
Converted 05c01 Loss functions.ipynb.
Converted 05c01a Training Utils for Flow Embedding.ipynb.
Converted 05c02 Testing MFE on Swiss Roll.ipynb.
Converted 05c03 Testing MFE on RNA Velocity Dataset.ipynb.
Converted 05c04 MFE with soft multiscale.ipynb.
Converted 05c05 Diffusion Distance based Loss - Tests and Visualizations-Copy1.ipynb.
Converted 05c05 Diffusion Distance based Loss - Tests and Visualizations.ipynb.
Converted 05c05b Learning flow around a fixed diffusion map.ipynb.
Converted 05d MFE All Pieces.ipynb.
Converted 05d01 MFE All Pieces-2.ipynb.
Converted 05d02 MFE Function Tests.ipynb.
Converted 06a Node2Vec Style Diffusion Flow Embedder.ipynb.
Converted 99k01 Playground.ipynb.
Converted Untitled.ipynb.
Converted index.ipynb.

Test: 2D Sinusoidal Ribbon

Test: 3D Sinusoidal Ribbon

from directed_graphs.datasets import plot_ribbon, sample_ribbon, plot_ribbon_samples
from directed_graphs.datasets import affinity_grid_search

Canceled future for execute_request message before replies were done

The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details.

plot_ribbon()
X, X_ = sample_ribbon()
labels = X
plot_ribbon_samples(X, X_)

X = torch.tensor(X)
X_ = torch.tensor(X_)
X = X.float().to(device)
X_ = X_.float().to(device)

affinity_grid_search(X, X_, sigmas=[0.5, 1, 3, 5], flow_strengths=[1, 2, 5, 10])

dfe = DiffusionFlowEmbedder(
    X,
    X_,
    t=1,
    sigma_graph=15,
    sigma_embedding=15,
    device=device,
    autoencoder_shape=[50, 10],
    flow_artist_shape=[30, 20, 10],
    flow_strength_graph=2,
    flow_strength_embedding=5,
    learnable_flow_strength=True,
    weight_of_flow=0.5,
    learning_rate=0.001,
)
dfe = dfe.to(device)
embeddings = dfe.fit(n_steps=2000)

dfe.visualize_points(labels)

Old Visualization Functions

Now migrated to "05c01a Flow Embedder Trainer"

#         # controls the x and y axes of the plot
#         # linspace(min on axis, max on axis, spacing on plot -- large number = more field arrows)
#         if labels is None:
#             labels = self.labels
#         uv = self.FlowArtist(self.grid).detach().cpu()
#         u = uv[:, 0].cpu()
#         v = uv[:, 1].cpu()
#         x = self.grid.detach().cpu()[:, 0]
#         y = self.grid.detach().cpu()[:, 1]
#         # quiver
#         #      plots a 2D field of arrows
#         #      quiver([X, Y], U, V, [C], **kw);
#         #      X, Y define the arrow locations, U, V define the arrow directions, and C optionally sets the color.
#         if labels is not None:
#             sc = plt.scatter(
#                 self.embedded_points[:, 0].detach().cpu(),
#                 self.embedded_points[:, 1].detach().cpu(),
#                 c=labels,
#             )
#             plt.legend()
#         else:
#             sc = plt.scatter(
#                 self.embedded_points[:, 0].detach().cpu(),
#                 self.embedded_points[:, 1].detach().cpu(),
#             )
#         plt.suptitle("Flow Embedding")
#         plt.quiver(x, y, u, v)
#         # Display all open figures.
#         plt.show()

#     def visualize_loss(self, loss_type="total"):
#         # diffusion_loss,reconstruction_loss, smoothness_loss
#         x = []
#         k = ""
#         losses = {}
#         for key in self.losses.keys():
#             losses[key] = []
#             k = key
#         losses["total"] = []
#         for i in range(len(self.losses["diffusion"])):
#             x.append(i)
#             for key in self.losses.keys():
#                 try:
#                     losses[key].append(self.losses[key][i].detach().cpu().numpy())
#                 except:
#                     losses[key].append(0)
#         if loss_type == "all":
#             for key in self.losses.keys():
#                 plt.plot(x, losses[key])
#             plt.legend(self.losses.keys(), loc="upper right")
#             plt.title("loss")
#         else:
#             plt.plot(x, losses[loss_type])
#             plt.title(loss_type)

#     def visualize_diffusion_matrices(self):
#         if len(self.ts) >= 3:
#             fig, axs = plt.subplots(3, 2, figsize=(10, 15))
#             axs[0][0].set_title(f"Ambient $P^{self.ts[0]}$")
#             axs[0][0].imshow(self.P_graph_ts[0].detach().cpu().numpy())
#             axs[0][1].set_title(f"Embedding $P^{self.ts[0]}$")
#             axs[0][1].imshow(self.P_embedding_ts[0].detach().cpu().numpy())

#             axs[1][0].set_title(f"Ambient $P^{self.ts[1]}$")
#             axs[1][0].imshow(self.P_graph_ts[1].detach().cpu().numpy())
#             axs[1][1].set_title(f"Embedding $P^{self.ts[1]}$")
#             axs[1][1].imshow(self.P_embedding_ts[1].detach().cpu().numpy())

#             axs[2][0].set_title(f"Ambient $P^{self.ts[2]}$")
#             axs[2][0].imshow(self.P_graph_ts[2].detach().cpu().numpy())
#             axs[2][1].set_title(f"Embedding $P^{self.ts[2]}$")
#             axs[2][1].imshow(self.P_embedding_ts[2].detach().cpu().numpy())
#         else:
#             fig, axs = plt.subplots(1, 2, figsize=(10, 5))
#             axs[0].set_title(f"Ambient $P^{self.ts[0]}$")
#             axs[0].imshow(self.P_graph_ts[0].detach().cpu().numpy())
#             axs[1].set_title(f"Embedding $P^{self.ts[0]}$")
#             axs[1].imshow(self.P_embedding_ts[0].detach().cpu().numpy())

#         plt.show()

`compute_grid`[source]

`diffusion_matrix_with_grid_points`[source]

General embedding matrix

`class` `FeedForwardReLU`[source]

The Flow Embedder

`class` `MultiscaleDiffusionFlowEmbedder`[source]

Basic Tests/Sanity Checks

Test: 2D Sinusoidal Ribbon

Test: 3D Sinusoidal Ribbon

Old Visualization Functions

compute_grid[source]

diffusion_matrix_with_grid_points[source]

General embedding matrix

class FeedForwardReLU[source]

The Flow Embedder

class MultiscaleDiffusionFlowEmbedder[source]

Basic Tests/Sanity Checks

Test: 2D Sinusoidal Ribbon

Test: 3D Sinusoidal Ribbon

Old Visualization Functions

`compute_grid`[source]

`diffusion_matrix_with_grid_points`[source]

`class` `FeedForwardReLU`[source]

`class` `MultiscaleDiffusionFlowEmbedder`[source]